//! program structure compiler
use std::{iter::Peekable, slice::Iter};

use crate::{
    compiler::{
        data::{
            SubroutineType, SymbolTable, SymbolTableItem, VariableDesc, VariableKind, VariableType,
            VariableTypePair,
        },
        vm_writer::VmWriter,
    },
    data::{Identifier, KeywordType, SymbolType},
    error::{CompilerError, CompilerResult},
    token::Token,
};

pub(super) struct ProgramCompilerUnit {
    /// the class symbol table
    class_table: SymbolTable,
    /// the subroutine symbol table
    subroutine_table: SymbolTable,
    /// the class name
    class_name: String,
    /// label counter
    label_counter: u16,
    /// vm code writer
    pub(super) vm_writer: VmWriter,
}
impl ProgramCompilerUnit {
    fn new() -> Self {
        Self {
            class_table: SymbolTable::new_tables(),
            subroutine_table: SymbolTable::new_tables(),
            class_name: String::default(),
            vm_writer: VmWriter::default(),
            label_counter: 0,
        }
    }
    pub(super) fn get_label_counter(&mut self) -> String {
        self.label_counter += 1;
        format!("{}.{}", self.class_name, self.label_counter)
    }
    /// generate general subroutine name
    pub(super) fn general_subroutine_name(&self, subroutine_name: &Identifier) -> String {
        format!("{}.{}", self.class_name, subroutine_name.as_str()) //the name rule: classname.functionname
    }
    fn reset_subroutine_table(&mut self) {
        self.subroutine_table = SymbolTable::new_tables();
    }
    /// get the variable info
    pub(super) fn get_variable(&self, var_name: &str) -> Option<&SymbolTableItem> {
        self.subroutine_table
            .get_row(var_name)
            .or_else(|| self.class_table.get_row(var_name))
    }
    /// The compiler flow unit of class
    pub(super) fn compile_class(tokens: &mut Peekable<Iter<'_, Token>>) -> CompilerResult<String> {
        let mut compiler = Self::new();
        // class
        let _ = crate::analyzer::Analyzer::match_keyword(
            tokens
                .next()
                .ok_or_else(|| CompilerError::Program("must start with class ;".into()))?,
            KeywordType::Class,
        )?;
        // calssName
        let class_name = super::expression::ExpressionCompilerUnit::flow_unit_identifier(
            tokens
                .next()
                .ok_or_else(|| CompilerError::Program("miss className".into()))?,
        )?;
        compiler.class_name = class_name.as_str().into();

        // {
        let _ = crate::analyzer::Analyzer::match_symbol(
            tokens
                .next()
                .ok_or_else(|| CompilerError::Program("miss {".into()))?,
            SymbolType::OpenBrace,
        )?;
        // classVarDec*
        while let Some(Token::Keyword(keyword)) = tokens.peek()
            && [KeywordType::Static, KeywordType::Field].contains(keyword)
        {
            compiler.compile_class_var_dec(tokens)?;
        }
        // subroutineDec*
        while let Some(Token::Keyword(keyword)) = tokens.peek()
            && [
                KeywordType::Constructor,
                KeywordType::Function,
                KeywordType::Method,
            ]
            .contains(keyword)
        {
            compiler.compile_subroutine_dec(tokens)?;
        }
        // }
        let _ = crate::analyzer::Analyzer::match_symbol(
            tokens
                .next()
                .ok_or_else(|| CompilerError::Program("miss }".into()))?,
            SymbolType::CloseBrace,
        )?;
        Ok(compiler.vm_writer.vm_code())
    }
    /// The compiler flow unit of class var decscript
    fn compile_class_var_dec(
        &mut self,
        tokens: &mut Peekable<Iter<'_, Token>>,
    ) -> CompilerResult<()> {
        // static | field
        if let Some(Token::Keyword(keyword)) = tokens.next()
            && [KeywordType::Static, KeywordType::Field].contains(keyword)
        {
            self.class_table.add_var_desc(
                (*keyword)
                    .try_into()
                    .expect("convert keyword to variable kind"),
                Self::compile_type_with_one_and_more_var_name(tokens)?,
            )?;

            let _ = crate::analyzer::Analyzer::match_symbol(
                tokens
                    .next()
                    .ok_or_else(|| CompilerError::Program("must end with ;".into()))?,
                SymbolType::Semicolon,
            )?;

            Ok(())
        } else {
            Err(CompilerError::Program(
                "class var dec must start with static or field".into(),
            ))
        }
    }

    /// The compiler flow unit of subroutine dec
    fn compile_subroutine_dec(
        &mut self,
        tokens: &mut Peekable<Iter<'_, Token>>,
    ) -> CompilerResult<()> {
        // constructor | function | method
        if let Some(Token::Keyword(keyword)) = tokens.next()
            && [
                KeywordType::Constructor,
                KeywordType::Function,
                KeywordType::Method,
            ]
            .contains(keyword)
        {
            let subroutine_type =
                SubroutineType::try_from(*keyword).expect("convert into subroutine type");

            // init subroutine table
            self.reset_subroutine_table();
            if subroutine_type == SubroutineType::Method {
                // add this to arg
                self.subroutine_table.add_row(
                    "this".into(),
                    VariableType::Class(self.class_name.clone()),
                    VariableKind::Arg,
                )?;
            }
            // void | type
            match tokens.peek() {
                Some(Token::Keyword(keyword)) if *keyword == KeywordType::Void => {
                    // skip void
                    let _ = tokens.next();
                }
                _ => {
                    let _ = Self::compile_type(tokens, false)?;
                }
            }
            // subroutine name
            let subroutine_name = super::expression::ExpressionCompilerUnit::flow_unit_identifier(
                tokens
                    .next()
                    .ok_or_else(|| CompilerError::Program("miss subroutine name".into()))?,
            )?;

            // (
            let _ = crate::analyzer::Analyzer::match_symbol(
                tokens
                    .next()
                    .ok_or_else(|| CompilerError::Program("miss parameterlist start (".into()))?,
                SymbolType::OpenParenthesis,
            )?;

            // parameterlist

            for pair in Self::compile_parameter_list(tokens)? {
                self.subroutine_table
                    .add_row(pair.var_name, pair.var_type, VariableKind::Arg)?;
            }
            // )
            let _ = crate::analyzer::Analyzer::match_symbol(
                tokens
                    .next()
                    .ok_or_else(|| CompilerError::Program("miss parameterlist end )".into()))?,
                SymbolType::CloseParenthesis,
            )?;

            // clear subroutine cache
            self.vm_writer.init_subroutine_cache();

            self.compile_subroutine_body(tokens)?;

            // insert function desc at the first
            // and write the subroutine cache to code.
            self.vm_writer.write_function(
                self.general_subroutine_name(&subroutine_name),
                self.subroutine_table.get_local_variable_num(),
                subroutine_type.clone(),
                self.class_table.get_field_variable_num(),
            );
        } else {
            return Err(CompilerError::Program(
                "subroutine must start with constructor function method or void".into(),
            ));
        }
        Ok(())
    }
    /// The compiler flow unit of parameter list
    fn compile_parameter_list(
        tokens: &mut Peekable<Iter<'_, Token>>,
    ) -> CompilerResult<Vec<VariableTypePair>> {
        let mut var_type_pairs = vec![];
        // type
        if Self::compile_type(tokens, true).is_ok() {
            // has parameter desc
            // type varName
            var_type_pairs.push(Self::compile_type_var_name(tokens)?);
            // (,type varName)*
            while let Some(Token::Symbol(SymbolType::Comma)) = tokens.peek() {
                // skip ,
                let _ = tokens.next();
                // type varName
                var_type_pairs.push(Self::compile_type_var_name(tokens)?);
            }
        }
        Ok(var_type_pairs)
    }

    /// The compiler flow unit of subroutine body
    fn compile_subroutine_body(
        &mut self,
        tokens: &mut Peekable<Iter<'_, Token>>,
    ) -> CompilerResult<()> {
        // {
        let _ = crate::analyzer::Analyzer::match_symbol(
            tokens.next().ok_or_else(|| {
                CompilerError::Program("subroutine body must start with {".into())
            })?,
            SymbolType::OpenBrace,
        )?;
        // varDec*
        while let Some(Token::Keyword(KeywordType::Var)) = tokens.peek() {
            self.subroutine_table
                .add_var_desc(VariableKind::Var, Self::compile_var_dec(tokens)?)?;
        }
        // statements
        super::statement::StatementCompilerUnit::compile_statements(self, tokens)?;
        // }
        let _ = crate::analyzer::Analyzer::match_symbol(
            tokens
                .next()
                .ok_or_else(|| CompilerError::Program("subroutine body must end with }".into()))?,
            SymbolType::CloseBrace,
        )?;

        Ok(())
    }

    /// The compiler flow unit of varDec
    fn compile_var_dec(tokens: &mut Peekable<Iter<'_, Token>>) -> CompilerResult<VariableDesc> {
        // var
        let _ = crate::analyzer::Analyzer::match_keyword(
            tokens
                .next()
                .ok_or_else(|| CompilerError::Program("must start with var".into()))?,
            KeywordType::Var,
        )?;
        let var_desc = Self::compile_type_with_one_and_more_var_name(tokens)?;
        let _ = crate::analyzer::Analyzer::match_symbol(
            tokens
                .next()
                .ok_or_else(|| CompilerError::Program("must end with ;".into()))?,
            SymbolType::Semicolon,
        )?;
        Ok(var_desc)
    }

    /// The compiler flow unit of type varName (, varName)*
    fn compile_type_with_one_and_more_var_name(
        tokens: &mut Peekable<Iter<'_, Token>>,
    ) -> CompilerResult<VariableDesc> {
        // type varName
        let VariableTypePair { var_type, var_name } = Self::compile_type_var_name(tokens)?;
        let mut vars = vec![var_name.as_str().to_string()];
        // , varName
        while let Some(Token::Symbol(SymbolType::Comma)) = tokens.peek() {
            let _ = tokens.next();
            vars.push(
                super::expression::ExpressionCompilerUnit::flow_unit_identifier(
                    tokens
                        .next()
                        .ok_or_else(|| CompilerError::Program("miss varName".into()))?,
                )?
                .as_str()
                .into(),
            );
        }
        Ok(VariableDesc { var_type, vars })
    }

    /// The compiler flow unit of type varName
    fn compile_type_var_name(
        tokens: &mut Peekable<Iter<'_, Token>>,
    ) -> CompilerResult<VariableTypePair> {
        // type
        let variable_type = Self::compile_type(tokens, false)?;
        // varName
        let variable_name = super::expression::ExpressionCompilerUnit::flow_unit_identifier(
            tokens
                .next()
                .ok_or_else(|| CompilerError::Program("miss varName".into()))?,
        )?;

        Ok(VariableTypePair::new(
            variable_type,
            variable_name.as_str().into(),
        ))
    }

    /// The compiler flow unit of type
    fn compile_type(
        tokens: &mut Peekable<Iter<'_, Token>>,
        is_peek: bool,
    ) -> CompilerResult<VariableType> {
        // int | char | boolean
        let token = if is_peek {
            tokens.peek().copied()
        } else {
            tokens.next()
        }
        .ok_or_else(|| CompilerError::Program("no token".into()))?;
        if let Token::Keyword(keyword) = token
            && [KeywordType::Int, KeywordType::Char, KeywordType::Boolean].contains(keyword)
        {
            return Ok((*keyword)
                .try_into()
                .expect("convert keyword to variable type"));
        }
        // className
        let identifier = super::expression::ExpressionCompilerUnit::flow_unit_identifier(token)?;
        Ok(VariableType::Class(identifier.as_str().into()))
    }
}
